
#include "copysign_tiling.h"
#include "register/op_def_registry.h"
#include <algorithm>
using namespace std;


namespace optiling {
static ge::graphStatus TilingFunc(gert::TilingContext* context)
{

    CopysignTilingData tiling;
        
    // 获取输入tensor的形状
    auto shape_x1 = context->GetInputShape(0)->GetOriginShape();
    // 有多少个维度
    int dimNum1 = shape_x1.GetDimNum();
    int totalLength1 = 1;
    int dim1[4] = {1, 1, 1, 1};
    // 计算每个维度的大小
    for (int i = 0; i < dimNum1; i++) {
        dim1[i] = shape_x1.GetDim(i);
        totalLength1 *= dim1[i];
    }
    // tiling.set_totalLength1(totalLength1);

    // 获取输入tensor的形状
    auto shape_x2 = context->GetInputShape(1)->GetOriginShape();
    // 有多少个维度
    int dimNum2 = shape_x2.GetDimNum();
    int totalLength2 = 1;
    int dim2[4] = {1, 1, 1, 1};
    // 计算每个维度的大小
    for (int i = 0; i < dimNum2; i++) {
        dim2[i] = shape_x2.GetDim(i);
        totalLength2 *= dim2[i];
    }
    // tiling.set_totalLength2(totalLength2);

    // tiling.set_dimNum1(dimNum1);
    // tiling.set_dimNum2(dimNum2);
    // tiling.set_dim1(dim1);
    tiling.set_dim2(dim2);

    int totalLength = 1;
    int dimNum = max(dimNum1, dimNum2);
    int dim[4] = {1, 1, 1, 1};
    for (int i = 0; i < dimNum; i++) {
        if (i < dimNum1) {
            dim[i] = max(dim[i], dim1[i]);
        }
        if (i < dimNum2) {
            dim[i] = max(dim[i], dim2[i]);
        }
        // printf("dim[%d] = %d\n", i, dim[i]);
        totalLength *= dim[i];
    }
    // printf("totalLength = %d\n", totalLength);
    // tiling.set_totalLength(totalLength);
    tiling.set_dim(dim);
    // tiling.set_dimNum(dimNum);
    
    int blockNum = 40;
    int blockLength = (totalLength2 + blockNum - 1) / blockNum;
    int tailBlockNum = totalLength2 % blockNum;

    // 特殊点
    if (totalLength == totalLength2) {
        context->SetTilingKey(1);
        if (totalLength <= 8192) {
            context->SetBlockDim(1);
        } else {
            context->SetBlockDim(40);
        }
    } else if (totalLength != totalLength2) {
        context->SetTilingKey(2);
        context->SetBlockDim(40);
    } else {
        context->SetTilingKey(3);
        // 每个核处理最后一维
        context->SetBlockDim(totalLength / dim[dimNum - 1]);
    }

    tiling.SaveToBuffer(context->GetRawTilingData()->GetData(), context->GetRawTilingData()->GetCapacity());
    context->GetRawTilingData()->SetDataSize(tiling.GetDataSize());
    return ge::GRAPH_SUCCESS;
}
}


namespace ge {
static ge::graphStatus InferShape(gert::InferShapeContext* context)
{
    const gert::Shape* x1_shape = context->GetInputShape(0);
    gert::Shape* y_shape = context->GetOutputShape(0);
    *y_shape = *x1_shape;
    return GRAPH_SUCCESS;
}
static ge::graphStatus InferDataType(gert::InferDataTypeContext *context)
{
const auto inputDataType = context->GetInputDataType(0);
context->SetOutputDataType(0, inputDataType);
return ge::GRAPH_SUCCESS;
}
}


namespace ops {
class Copysign : public OpDef {
public:
    explicit Copysign(const char* name) : OpDef(name)
    {
        this->Input("input")
            .ParamType(REQUIRED)
            .DataType({ge::DT_FLOAT, ge::DT_FLOAT16, ge::DT_BF16})
            .Format({ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND})
            .UnknownShapeFormat({ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND});
        this->Input("other")
            .ParamType(REQUIRED)
            .DataType({ge::DT_FLOAT, ge::DT_FLOAT16, ge::DT_BF16})
            .Format({ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND})
            .UnknownShapeFormat({ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND});
        this->Output("out")
            .ParamType(REQUIRED)
            .DataType({ge::DT_FLOAT, ge::DT_FLOAT16, ge::DT_BF16})
            .Format({ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND})
            .UnknownShapeFormat({ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND});

        this->SetInferShape(ge::InferShape).SetInferDataType(ge::InferDataType);

        this->AICore()
            .SetTiling(optiling::TilingFunc);
        this->AICore().AddConfig("ascend910b");

    }
};

OP_ADD(Copysign);
}
